from BPRMF.Recommender import Recommender as BPRRecommender
from POP.Recommender import Recommender as POPRecommender
from ICF.Recommender import Recommender as ICFRecommender
from UCF.Recommender import Recommender as UCFRecommender
from typing import Dict, List
from torch.cuda import is_available
from tqdm import tqdm
import numpy as np
import time


class Recommender:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        self.bpr_recommender = BPRRecommender(hyper_params)
        self.pop_recommender = POPRecommender(hyper_params)
        self.icf_recommender = ICFRecommender(hyper_params)
        self.ucf_recommender = UCFRecommender(hyper_params)

        # POP Recommender are used for filling the remain items
        self.total_cnt = hyper_params['total_cnt']
        self.bpr_cnt = hyper_params['bpr_cnt']
        self.icf_cnt = hyper_params['icf_cnt']
        self.ucf_cnt = hyper_params['ucf_cnt']

    def recommend(self) -> Dict[int, List[int]]:
        bpr_result = self.bpr_recommender.recommend(self.bpr_cnt)
        icf_result = self.icf_recommender.recommend(self.icf_cnt)
        ucf_result = self.ucf_recommender.recommend(self.ucf_cnt)
        pop_result = self.pop_recommender.recommend(self.total_cnt)

        result: Dict[int, List[int]] = {}
        user_cnt = self.pop_recommender.user_cnt

        print('Recommender: Merging results...')
        for user in tqdm(range(user_cnt)):
            cur_result = set()

            if user in bpr_result:
                cur_result = cur_result.union(set(bpr_result[user]))
            if user in icf_result:
                cur_result = cur_result.union(set(icf_result[user]))
            if user in ucf_result:
                cur_result = cur_result.union(set(ucf_result[user]))

            for rec in pop_result:
                cur_result.add(rec)

                if cur_result.__len__() >= self.total_cnt:
                    break

            result[user] = list(cur_result)
            if len(result[user]) > self.total_cnt:
                result[user] = np.random.choice(
                    result[user], self.total_cnt, replace=False)

        return result


if __name__ == '__main__':
    hyper_params = {
        'dataset_path': 'bslen15/dataset_out_2.txt',
        'bind_file_path': 'bslen15/',
        'load_from_file': False,
        'device': 'cuda' if is_available() else 'cpu',
        'batch_size': 200 if is_available() else 64,
        'epochs': 5 if is_available() else 2,
        'embed_dim': 64,
        'neg_sample_cnt': 300,
        'epsilon': 0.1,
        'total_cnt': 500,
        'bpr_cnt': 200,
        'icf_cnt': 100,
        'ucf_cnt': 100
    }

    print('Recommender: Initializing...')
    recommender = Recommender(hyper_params)
    print('Recommender: Finished.')

    start_time = time.time()
    print('Recommender: Start recommend...')
    result = recommender.recommend()
    print(f'Recommender: Finished. Time: {time.time()-start_time} seconds.')
    print(len(result[0]))
